"""
YOLO目标检测模块
集成YOLOv11进行目标检测，支持多种目标类型识别
"""

import cv2
import numpy as np
from ultralytics import YOLO
from typing import List, Dict, Tuple, Optional
import logging
import os

class YOLODetector:
    """YOLO目标检测器"""
    
    def __init__(self, model_path: str = "yolo11n.pt", confidence: float = 0.5):
        """
        初始化YOLO检测器
        
        Args:
            model_path: YOLO模型路径
            confidence: 置信度阈值
        """
        self.model_path = model_path
        self.confidence = confidence
        self.model: Optional[YOLO] = None
        self.target_model: Optional[YOLO] = None  # 专门用于靶子检测的模型
        
        # 目标类别配置
        self.target_classes = {
            'person': 0,  # 人员（荧光衣检测）
            'target': 0,  # 靶子
            'bullseye': 1,  # 靶心
            'ring': 2,  # 环数
        }
        
        # 目标颜色配置
        self.class_colors = {
            'person': (0, 255, 0),      # 绿色 - 人员
            'target': (0, 0, 255),      # 红色 - 靶子
            'bullseye': (255, 0, 0),    # 蓝色 - 靶心
            'ring': (0, 255, 255),      # 青色 - 环数
            'default': (255, 255, 0),   # 黄色 - 默认
        }
        
        # 设置日志
        self.logger = logging.getLogger(__name__)
        
    def initialize(self) -> bool:
        """
        初始化YOLO模型
        
        Returns:
            bool: 初始化是否成功
        """
        try:
            # 检查模型文件是否存在
            if not os.path.exists(self.model_path):
                self.logger.error(f"模型文件不存在: {self.model_path}")
                return False
            
            # 加载YOLO模型
            self.model = YOLO(self.model_path)
            self.logger.info(f"YOLO模型加载成功: {self.model_path}")
            
            # 获取模型类别信息
            if hasattr(self.model, 'names'):
                self.logger.info(f"模型支持的类别: {self.model.names}")
            
            # 尝试加载靶子检测模型
            self.load_target_model()
            
            return True
            
        except Exception as e:
            self.logger.error(f"YOLO模型初始化失败: {e}")
            return False
    
    def load_target_model(self, target_model_path: Optional[str] = None) -> bool:
        """
        加载靶子检测模型
        
        Args:
            target_model_path: 靶子模型路径，如果为None则使用默认路径
            
        Returns:
            bool: 加载是否成功
        """
        try:
            # 默认靶子模型路径
            if target_model_path is None:
                target_model_path = "models/target_detection.pt"
            
            # 检查模型文件是否存在
            if not os.path.exists(target_model_path):
                self.logger.warning(f"靶子检测模型不存在: {target_model_path}")
                self.logger.info("将使用通用YOLO模型进行靶子检测")
                return False
            
            # 加载靶子检测模型
            self.target_model = YOLO(target_model_path)
            self.logger.info(f"靶子检测模型加载成功: {target_model_path}")
            
            # 获取靶子模型类别信息
            if hasattr(self.target_model, 'names'):
                self.logger.info(f"靶子模型支持的类别: {self.target_model.names}")
            
            return True
            
        except Exception as e:
            self.logger.error(f"靶子检测模型加载失败: {e}")
            return False
    
    def train_target_model(self, dataset_yaml: str = "datasets/targets/target.yaml", 
                          epochs: int = 100, imgsz: int = 640) -> bool:
        """
        训练靶子检测模型
        
        Args:
            dataset_yaml: 数据集配置文件路径
            epochs: 训练轮数
            imgsz: 图像尺寸
            
        Returns:
            bool: 训练是否成功
        """
        try:
            # 检查数据集配置文件是否存在
            if not os.path.exists(dataset_yaml):
                self.logger.error(f"数据集配置文件不存在: {dataset_yaml}")
                return False
            
            # 创建模型目录
            os.makedirs("models", exist_ok=True)
            
            # 使用YOLOv11n作为基础模型进行训练
            model = YOLO("yolo11n.pt")
            
            self.logger.info(f"开始训练靶子检测模型...")
            self.logger.info(f"数据集: {dataset_yaml}")
            self.logger.info(f"训练轮数: {epochs}")
            self.logger.info(f"图像尺寸: {imgsz}")
            
            # 开始训练
            results = model.train(
                data=dataset_yaml,
                epochs=epochs,
                imgsz=imgsz,
                project="models",
                name="target_detection",
                save=True,
                verbose=True
            )
            
            # 保存最佳模型
            best_model_path = "models/target_detection/weights/best.pt"
            if os.path.exists(best_model_path):
                # 复制到标准位置
                import shutil
                shutil.copy2(best_model_path, "models/target_detection.pt")
                self.logger.info(f"靶子检测模型训练完成: models/target_detection.pt")
                
                # 加载新训练的模型
                self.load_target_model("models/target_detection.pt")
                return True
            else:
                self.logger.error("训练完成但未找到最佳模型文件")
                return False
                
        except Exception as e:
            self.logger.error(f"靶子检测模型训练失败: {e}")
            return False
    
    def detect(self, frame: np.ndarray) -> List[Dict]:
        """
        执行目标检测
        
        Args:
            frame: 输入图像帧
            
        Returns:
            List[Dict]: 检测结果列表，每个结果包含bbox、confidence、class等信息
        """
        if self.model is None:
            self.logger.error("YOLO模型未初始化")
            return []
        
        try:
            # 执行检测
            results = self.model(frame, conf=self.confidence, verbose=False)
            
            detections = []
            for result in results:
                boxes = result.boxes
                if boxes is not None:
                    for box in boxes:
                        # 获取边界框坐标
                        x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
                        
                        # 获取置信度和类别
                        confidence = float(box.conf[0].cpu().numpy())
                        class_id = int(box.cls[0].cpu().numpy())
                        
                        # 获取类别名称
                        class_name = self.model.names.get(class_id, f"class_{class_id}")
                        
                        detection = {
                            'bbox': [int(x1), int(y1), int(x2), int(y2)],
                            'confidence': confidence,
                            'class_id': class_id,
                            'class_name': class_name,
                            'center': [int((x1 + x2) / 2), int((y1 + y2) / 2)],
                            'area': int((x2 - x1) * (y2 - y1))
                        }
                        
                        detections.append(detection)
            
            return detections
            
        except Exception as e:
            self.logger.error(f"目标检测失败: {e}")
            return []
    
    def detect_targets(self, frame: np.ndarray) -> List[Dict]:
        """
        专门检测靶子的方法
        
        Args:
            frame: 输入图像帧
            
        Returns:
            List[Dict]: 靶子检测结果列表
        """
        # 如果有专门的靶子模型，使用靶子模型
        if self.target_model is not None:
            try:
                results = self.target_model(frame, conf=self.confidence, verbose=False)
                
                detections = []
                for result in results:
                    boxes = result.boxes
                    if boxes is not None:
                        for box in boxes:
                            # 获取边界框坐标
                            x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
                            
                            # 获取置信度和类别
                            confidence = float(box.conf[0].cpu().numpy())
                            class_id = int(box.cls[0].cpu().numpy())
                            
                            # 获取类别名称
                            class_name = self.target_model.names.get(class_id, f"class_{class_id}")
                            
                            detection = {
                                'bbox': [int(x1), int(y1), int(x2), int(y2)],
                                'confidence': confidence,
                                'class_id': class_id,
                                'class_name': class_name,
                                'center': [int((x1 + x2) / 2), int((y1 + y2) / 2)],
                                'area': int((x2 - x1) * (y2 - y1)),
                                'type': 'target'  # 标记为靶子检测结果
                            }
                            
                            detections.append(detection)
                
                return detections
                
            except Exception as e:
                self.logger.error(f"靶子检测失败: {e}")
                return []
        else:
            # 如果没有专门的靶子模型，使用通用模型
            self.logger.warning("使用通用模型进行靶子检测（精度可能较低）")
            return self.detect(frame)
    
    def draw_detections(self, frame: np.ndarray, detections: List[Dict]) -> np.ndarray:
        """
        在图像上绘制检测结果
        
        Args:
            frame: 输入图像帧
            detections: 检测结果列表
            
        Returns:
            np.ndarray: 绘制了检测结果的图像
        """
        result_frame = frame.copy()
        
        for detection in detections:
            bbox = detection['bbox']
            confidence = detection['confidence']
            class_name = detection['class_name']
            center = detection['center']
            
            # 获取颜色
            color = self.class_colors.get(class_name, self.class_colors['default'])
            
            # 绘制边界框
            cv2.rectangle(result_frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color, 2)
            
            # 绘制中心点
            cv2.circle(result_frame, tuple(center), 5, color, -1)
            
            # 绘制标签
            label = f"{class_name}: {confidence:.2f}"
            label_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)[0]
            
            # 标签背景
            cv2.rectangle(result_frame, 
                         (bbox[0], bbox[1] - label_size[1] - 10),
                         (bbox[0] + label_size[0], bbox[1]),
                         color, -1)
            
            # 标签文字
            cv2.putText(result_frame, label,
                       (bbox[0], bbox[1] - 5),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
        
        return result_frame
    
    def filter_detections(self, detections: List[Dict], 
                         target_classes: Optional[List[str]] = None,
                         min_confidence: Optional[float] = None,
                         min_area: Optional[int] = None) -> List[Dict]:
        """
        过滤检测结果
        
        Args:
            detections: 原始检测结果
            target_classes: 目标类别列表
            min_confidence: 最小置信度
            min_area: 最小面积
            
        Returns:
            List[Dict]: 过滤后的检测结果
        """
        filtered = detections.copy()
        
        # 按类别过滤
        if target_classes:
            filtered = [d for d in filtered if d['class_name'] in target_classes]
        
        # 按置信度过滤
        if min_confidence is not None:
            filtered = [d for d in filtered if d['confidence'] >= min_confidence]
        
        # 按面积过滤
        if min_area is not None:
            filtered = [d for d in filtered if d['area'] >= min_area]
        
        return filtered
    
    def get_best_target(self, detections: List[Dict], 
                       prefer_class: Optional[str] = None) -> Optional[Dict]:
        """
        获取最佳目标（用于瞄准）
        
        Args:
            detections: 检测结果列表
            prefer_class: 优先类别
            
        Returns:
            Optional[Dict]: 最佳目标，如果没有则返回None
        """
        if not detections:
            return None
        
        # 如果指定了优先类别，先尝试找该类别
        if prefer_class:
            preferred = [d for d in detections if d['class_name'] == prefer_class]
            if preferred:
                # 返回置信度最高的
                return max(preferred, key=lambda x: x['confidence'])
        
        # 否则返回置信度最高的目标
        return max(detections, key=lambda x: x['confidence'])
    
    def set_confidence_threshold(self, confidence: float):
        """
        设置置信度阈值
        
        Args:
            confidence: 新的置信度阈值
        """
        self.confidence = max(0.0, min(1.0, confidence))
        self.logger.info(f"置信度阈值已设置为: {self.confidence}")
    
    def get_model_info(self) -> Dict:
        """
        获取模型信息
        
        Returns:
            Dict: 模型信息
        """
        if self.model is None:
            return {}
        
        info = {
            'model_path': self.model_path,
            'confidence': self.confidence,
            'classes': getattr(self.model, 'names', {}),
            'target_classes': self.target_classes,
        }
        
        return info